/*
 * Copyright 2019-2025 the original author or authors.
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 * https://www.apache.org/licenses/LICENSE-2.0
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package com.iiifi.kite.rocketmq.configuration;

import java.util.HashMap;
import java.util.Map;
import java.util.Properties;
import java.util.UUID;

import javax.annotation.PostConstruct;

import org.apache.commons.lang3.StringUtils;
import org.apache.rocketmq.client.consumer.DefaultMQPushConsumer;
import org.apache.rocketmq.client.consumer.listener.MessageListenerConcurrently;
import org.apache.rocketmq.client.consumer.listener.MessageListenerOrderly;
import org.apache.rocketmq.client.exception.MQClientException;
import org.apache.rocketmq.common.protocol.heartbeat.MessageModel;
import org.springframework.boot.autoconfigure.condition.ConditionalOnBean;
import org.springframework.context.annotation.Configuration;
import org.springframework.core.env.Environment;
import org.springframework.util.Assert;
import org.springframework.util.CollectionUtils;

import com.iiifi.kite.rocketmq.annotation.RocketMqConsumer;
import com.iiifi.kite.rocketmq.core.AbstractMqPushConsumer;
import com.iiifi.kite.rocketmq.core.MessageExtConst;
import com.iiifi.kite.rocketmq.trace.common.OnsTraceConstants;
import com.iiifi.kite.rocketmq.trace.dispatch.impl.AsyncTraceAppender;
import com.iiifi.kite.rocketmq.trace.dispatch.impl.AsyncTraceDispatcher;
import com.iiifi.kite.rocketmq.trace.tracehook.OnsConsumeMessageHookImpl;

import lombok.extern.slf4j.Slf4j;

/**
 * 自动装配消息消费者
 *
 * @author kite@iiifi.com 花朝
 */
@Slf4j
@Configuration
@ConditionalOnBean(RocketMqAutoConfiguration.class)
public class RocketMqConsumerAutoConfiguration extends RocketMqAutoConfiguration {

    private AsyncTraceDispatcher asyncTraceDispatcher;
    /**
     * 维护一份map用于检测是否用同样的consumerGroup订阅了不同的topic+tag
     */
    private Map<String, String> validConsumerMap;

    @PostConstruct
    public void init() throws Exception {
        Map<String, Object> beans = applicationContext.getBeansWithAnnotation(RocketMqConsumer.class);
        if (!CollectionUtils.isEmpty(beans) && mqProperties.getTraceEnabled()) {
            initAsyncAppender();
        }
        validConsumerMap = new HashMap<>(16);
        for (Map.Entry<String, Object> entry : beans.entrySet()) {
            publishConsumer(entry.getKey(), entry.getValue());
        }
        // 清空map，等待回收
        validConsumerMap = null;
    }

    private AsyncTraceDispatcher initAsyncAppender() {
        if (asyncTraceDispatcher != null) {
            return asyncTraceDispatcher;
        }
        try {
            Properties tempProperties = new Properties();
            tempProperties.put(OnsTraceConstants.MaxMsgSize, "128000");
            tempProperties.put(OnsTraceConstants.AsyncBufferSize, "2048");
            tempProperties.put(OnsTraceConstants.MaxBatchNum, "1");
            tempProperties.put(OnsTraceConstants.WakeUpNum, "1");
            tempProperties.put(OnsTraceConstants.NAMESRV_ADDR, mqProperties.getNameServerAddress());
            tempProperties.put(OnsTraceConstants.InstanceName, UUID.randomUUID().toString());
            AsyncTraceAppender asyncAppender = new AsyncTraceAppender(tempProperties);
            asyncTraceDispatcher = new AsyncTraceDispatcher(tempProperties);
            asyncTraceDispatcher.start(asyncAppender, "DEFAULT_WORKER_NAME");
        } catch (MQClientException e) {
            e.printStackTrace();
        }
        return asyncTraceDispatcher;
    }

    private void publishConsumer(String beanName, Object bean) throws Exception {
        RocketMqConsumer rocketMqConsumer = applicationContext.findAnnotationOnBean(beanName, RocketMqConsumer.class);
        if (StringUtils.isEmpty(mqProperties.getNameServerAddress())) {
            throw new RuntimeException("name server address must be defined");
        }
        assert rocketMqConsumer != null;
        Assert.notNull(rocketMqConsumer, "consumer must be defined");
        Assert.notNull(rocketMqConsumer.consumerGroup(), "consumer's consumerGroup must be defined");
        Assert.notNull(rocketMqConsumer.topic(), "consumer's topic must be defined");
        if (!AbstractMqPushConsumer.class.isAssignableFrom(bean.getClass())) {
            throw new RuntimeException(bean.getClass().getName() + " - consumer未实现Consumer抽象类");
        }
        Environment environment = applicationContext.getEnvironment();
        String consumerGroup = environment.resolvePlaceholders(rocketMqConsumer.consumerGroup());
        String topic = environment.resolvePlaceholders(rocketMqConsumer.topic());
        String tags = "*";
        if (rocketMqConsumer.tag().length == 1) {
            tags = environment.resolvePlaceholders(rocketMqConsumer.tag()[0]);
        } else if (rocketMqConsumer.tag().length > 1) {
            tags = StringUtils.join(rocketMqConsumer.tag(), "||");
        }
        // 检查consumerGroup
        if (!StringUtils.isEmpty(validConsumerMap.get(consumerGroup))) {
            String exist = validConsumerMap.get(consumerGroup);
            throw new RuntimeException("消费组重复订阅，请新增消费组用于新的topic和tag组合: " + consumerGroup + "已经订阅了" + exist);
        } else {
            validConsumerMap.put(consumerGroup, topic + "-" + tags);
        }
        // 配置push consumer
        DefaultMQPushConsumer consumer = new DefaultMQPushConsumer(consumerGroup);
        consumer.setNamesrvAddr(mqProperties.getNameServerAddress());
        consumer.setMessageModel(MessageModel.valueOf(rocketMqConsumer.messageMode()));
        consumer.subscribe(topic, tags);
        consumer.setInstanceName(UUID.randomUUID().toString());
        consumer.setVipChannelEnabled(mqProperties.getVipChannelEnabled());
        AbstractMqPushConsumer abstractMqPushConsumer = (AbstractMqPushConsumer) bean;
        if (MessageExtConst.CONSUME_MODE_CONCURRENTLY.equals(rocketMqConsumer.consumeMode())) {
            consumer.registerMessageListener((MessageListenerConcurrently) abstractMqPushConsumer::dealMessage);
        } else if (MessageExtConst.CONSUME_MODE_ORDERLY.equals(rocketMqConsumer.consumeMode())) {
            consumer.registerMessageListener((MessageListenerOrderly) abstractMqPushConsumer::dealMessage);
        } else {
            throw new RuntimeException("unknown consume mode ! only support CONCURRENTLY and ORDERLY");
        }
        abstractMqPushConsumer.setConsumer(consumer);
        // 为Consumer增加消息轨迹回发模块
        if (mqProperties.getTraceEnabled()) {
            try {
                consumer.getDefaultMQPushConsumerImpl().registerConsumeMessageHook(
                        new OnsConsumeMessageHookImpl(asyncTraceDispatcher));
            } catch (Throwable e) {
                log.error("system mqtrace hook init failed ,maybe can't send msg trace data");
            }
        }
        consumer.start();
        log.info(String.format("%s is ready to subscribe message", bean.getClass().getName()));
    }
}
